import numpy as np

from init import *
from sklearn.preprocessing import MinMaxScaler
import copy
from collections import Counter
from scipy.optimize import minimize
from init import *
import pandas as pd
# from ga_method import *
from gatest import *
import math

plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False


def get_feature(img_data, flag):
    '''
    对图像的特征进行提取
    :param img_data: 图像信息
    :param flag: 奇数点还是偶数点
    :return:
    '''
    row, col = img_data.shape
    size = int((row * col) / 2)
    errolist = np.zeros(size)
    em = np.zeros((row, col))
    sur = np.zeros(size)
    index = 0
    for i in range(2, row - 2):
        for j in range(2, col - 2):
            # # 像素上下左右四个像素
            v1 = int(img_data[i - 1][j])
            v2 = int(img_data[i][j - 1])
            v3 = int(img_data[i + 1][j])
            v4 = int(img_data[i][j + 1])
            p1 = (math.modf(np.average([v1, v2, v3, v4])))[0]
            p = np.ceil(np.average([v1, v2, v3, v4])) if p1 >= 0.5 else np.floor(np.average([v1, v2, v3, v4]))
            # 左梯度和右梯度
            surround = np.array([img_data[i, j - 1], img_data[i - 1, j], img_data[i, j + 1], img_data[i + 1, j]])
            # print(surround)
            if np.mod(i + j, 2) == 0 and flag == 0:
                # dist_near = np.diff(surround)
                sur[index] = np.var(surround,ddof=0)
                errolist[index] = int(img_data[i][j]) - int(p)
                em[i][j] = int(img_data[i][j]) - int(p)
                index += 1
            if np.mod(i + j, 2) == 1 and flag == 1:
                # dist_near = np.diff(surround)
                sur[index] = np.var(surround,ddof=0)
                errolist[index] = int(img_data[i][j]) - int(p)
                em[i][j] = int(img_data[i][j]) - int(p)
                # print(em[i][j])
                index += 1
    sur = sur[:index]

    sur = sur.reshape(-1,1)
    errolist = errolist[:index]
    scaler = MinMaxScaler()#feature_range=(0, 1000)
    sur = scaler.fit_transform(sur)
    sur = sur.flatten()
    if flag == 0:
        file_names = 'sur3.mat'
        scio.savemat(file_names, {'susss': sur})
    # sizes = sur.shape[0]
    # sur = sur.reshape(sizes,-1)
    print(sur.shape)
    idx = np.argsort(sur)
    errolist = errolist[idx]
    return errolist, index, em






def complexity_map(img_data, label, flags):
    cm = np.zeros(img_data.shape)
    A, B = img_data.shape
    count = 0
    for i in range(2, A - 2):
        for j in range(2, B - 2):
            if np.mod(i + j, 2) == 0 and flags == 0:
                cm[i][j] = int(label[count]) + 1
                count = count + 1
            if np.mod(i + j, 2) == 1 and flags == 1:
                cm[i][j] = int(label[count]) + 1
                count = count + 1
    return cm


def erro_map(img_data, flag2):
    row, col = img_data.shape
    erro = np.zeros(img_data.shape)
    erro_list = []
    for i in range(2, row - 2):
        for j in range(2, col - 2):
            v1 = int(img_data[i - 1][j])
            v2 = int(img_data[i][j - 1])
            v3 = int(img_data[i + 1][j])
            v4 = int(img_data[i][j + 1])
            p3 = np.ceil(np.average([v1, v2, v3, v4]))
            if np.mod(i + j, 2) == 0 and flag2 == 0:
                erro[i][j] = img_data[i][j] - p3
                erro_list.append(img_data[i][j] - p3)
            if np.mod(i + j, 2) == 1 and flag2 == 1:
                erro[i][j] = img_data[i][j] - p3
                erro_list.append(img_data[i][j] - p3)
    return erro, erro_list


def cal_entropy(imgs, label, ks, flags):
    A, B = imgs.shape
    his = []
    count = 0
    entropy = []
    for i in range(ks):
        his.append([])
    for i in range(2, A - 2):
        for j in range(2, B - 2):
            v1 = int(imgs[i - 1][j])
            v2 = int(imgs[i + 1][j])
            v3 = int(imgs[i][j + 1])
            v4 = int(imgs[i][j - 1])
            e = np.ceil(np.average([v1, v2, v3, v4])) - int(imgs[i][j])
            if flags == 0 and np.mod(i + j, 2) == 0:
                his[label[count]].append(e)
                count = count + 1
            elif flags == 1 and np.mod(i + j, 2) == 1:
                his[label[count]].append(e)
                count = count + 1

    for i in range(ks):
        e = 0
        c = Counter(his[i])
        total = len(his[i])
        for i in c.values():
            pro = int(i) / total
            e = e + pro * math.log(1 / pro, 2)
        entropy.append(e)
    print(entropy)
    entropy_sort = sorted(entropy);
    cor = []
    labels = []
    print(entropy_sort)
    for i in range(ks):
        for j in range(ks):
            if entropy[i] == entropy_sort[j]:
                cor.append(j)
                entropy_sort[j] = -1
                break
    print(cor)
    for i in label:
        labels.append(cor[i])
    legend = []
    fig = []
    frequency = list(range(ks))
    new_name = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p']
    for i in range(ks):
        legend.append('l=' + str(cor[i]))
        fig.append(str(cor[i]))
        # print(len(his[i]))
        frequency[cor[i]] = int(len(his[cor[i]]))
    plt.grid(ls=':');
    plt.bar(range(ks), frequency, tick_label=new_name, color='b');
    plt.title('The distribution of pixels');
    plt.xlabel('Neuron')
    plt.ylabel('Population');
    print(10086)
    plt.show()
    return labels, cor


def cal_aux(s, data_length,lcm):
    '''

        lcm是localtionmapc
    '''
    if np.sum(lcm) == 0:
        Aux = 18*len(s) + 4 + 1 + 17 + 4
    else:
        Aux = 18*len(s) + 4 + 1 + 17 + 4+8*data_length
    return Aux


def re_hs(errolist,em,flag):
    Hs = np.zeros((512,512))
    row,col  =em.shape
    index = 0
    for i in range(2, row - 2):
        for j in range(2, col - 2):
            if np.mod(i+j,2) == 0 and flag ==0:
                Hs[i][j] = errolist[index]
                index+=1
            if np.mod(i+j,2) == 1 and flag ==1:
                Hs[i][j] = errolist[index]
                index += 1
    return Hs


def insert_msg(em, cm, bin, img_hide3, secret, ec3, flag3):
    '''
        em误差矩阵
        Dbis复杂度矩阵
        bin最佳区间
        img_hide图形矩阵
        secret嵌入矩阵
        插入需要插入的信息
    '''
    index = 0
    index2 = 0
    row, col = img_hide3.shape
    for i in range(2, row - 2):
        for j in range(2, col - 2):
            if np.mod(i + j, 2) == 0 and flag3 == 0 and ec3 > 0:
                ks = cm[i][j] - 1
                if em[i][j] == bin[int(ks)]:
                    ec3 = ec3 - 1
                    img_hide3[i][j] += secret[i][j]
                    if secret[i][j] == 1:
                        index2 += 1
                if em[i][j] == -bin[int(ks)] - 1:
                    ec3 = ec3 - 1
                    img_hide3[i][j] -= secret[i][j]
                    if secret[i][j] == 1:
                        index2 += 1
                if em[i][j] > bin[int(ks)]:
                    img_hide3[i][j] += 1
                    index += 1
                if em[i][j] < -bin[int(ks)] - 1:
                    img_hide3[i][j] -= 1
                    index += 1
            if np.mod(i + j, 2) == 1 and flag3 == 1 and ec3 > 0:
                ks = cm[i][j] - 1
                if em[i][j] == bin[int(ks)]:
                    ec3 = ec3 - 1
                    img_hide3[i][j] += secret[i][j]
                    if secret[i][j] == 1:
                        index2 += 1
                if em[i][j] == -bin[int(ks)] - 1:
                    ec3 = ec3 - 1
                    img_hide3[i][j] -= secret[i][j]
                    if secret[i][j] == 1:
                        index2 += 1
                if em[i][j] > bin[int(ks)]:
                    img_hide3[i][j] += 1
                    index += 1
                if em[i][j] < -bin[int(ks)] - 1:
                    img_hide3[i][j] -= 1
                    index += 1

    print("偏移数量", index)
    print('偏移点数量', index2)
    return img_hide3, ec3


def danci(danci_img, flags, ecfinals):
    '''

    :param danci_img: 初始图像矩阵
    :param flags: x/o
    :param ecfinals: 嵌入容量
    :return: 获取特征，构造特征矩阵，归一化矩阵，进行ga算法求最优Bin，计算边信息进行嵌入
    '''
    # 获取图片的坐标的长宽
    row, col = danci_img.shape
    danci_img, lcm, halfsize, data = Over_flow(danci_img, flags)
    # 数组长度
    data_length = len(data)
    errolist, indexsize, em = get_feature(danci_img, flags)
    em2 = re_hs(errolist,em,flag)
    # 计算嵌入点
    ec1 =np.ceil(ecfinals / 2 + 20)
    P, Z, idx, Dist = devideAndPick2(errolist, ec1, flags)
    aux = cal_aux(P,data_length,lcm)
    secret = secret_msg(row, col)
    P,Z = getModifyMap(P,Z)
    imgs, ec1 = inster_ms(em2, danci_img, secret, ec1, flags, P, Z)
    anx_secret = secret_msg(row, col)
    imgs = Lsb_insert(row, col, aux, imgs, flags, anx_secret)
    return imgs,ec1


def devideAndPick2(errs, rate, flags):
    P = []
    Z = []
    dist_lists = np.ones(100) * 100000000
    pl_size = np.zeros((100, 2))
    tmpArea = np.zeros((100, 2))
    Dist = 10000000000
    miner = np.ceil(rate / len(errs) * 100)
    maxer = 99

    dist_lists, data = Ga_one(errs, miner, dist_lists, rate, flags)
    if dist_lists[int(miner)] != 0 and dist_lists[int(miner)] < Dist:
        Dist = dist_lists[int(miner)]
        P = data[0]
        Z = data[1]

    dist_lists, data = Ga_one(errs, maxer, dist_lists, rate, flags)
    if dist_lists[int(maxer)] != 0 and dist_lists[int(maxer)] < Dist:
        Dist = dist_lists[int(maxer)]
        P = data[0]
        Z = data[1]
    temp = 1
    pl_size[0, :] = [miner, maxer]
    if miner + 1 == maxer:
        temp = 0
    index = 0
    aa = []
    while temp > 0:
        tmpC = 0
        for i in range(temp):
            # print('plsize',pl_size)
            idx = np.floor((pl_size[i, 0] + pl_size[i, 1]) / 2)
            aa.append([idx, pl_size[i, 0], pl_size[i, 1]])
            print('这次的', idx, pl_size[i, 0], pl_size[i, 1])
            dist_lists, data = Ga_one(errs, idx, dist_lists, rate, flags)
            if dist_lists[int(idx)] != 0 and dist_lists[int(idx)] < Dist:
                Dist = dist_lists[int(idx)]
                P = data[0]
                Z = data[1]
            if idx > pl_size[i, 0] + 1:
                # print('tmc1',tmpC)
                tmpArea[tmpC, :] = [pl_size[i, 0], idx]
                tmpC += 1
            if idx + 1 < pl_size[i, 1]:
                print('tmc2', tmpC)
                tmpArea[tmpC, :] = [idx, pl_size[i, 1]]
                tmpC += 1
        temp = 0
        minindex = np.argmin(dist_lists)
        for i in range(tmpC):
            if tmpArea[i, 0] == minindex or tmpArea[i, 1] == minindex:
                pl_size[temp, :] = tmpArea[i, :]
                temp += 1
    idx = np.argmin(dist_lists)
    return P, Z, idx, dist_lists[idx]


def inster_ms(em, img_hide3, secret, ec3, flag3, P, Z):
    row, col = img_hide3.shape
    index = 0
    for i in range(2, row - 2):
        for j in range(2, col - 2):
            if np.mod(i + j, 2) == 0 and flag3 == 0 and ec3 > 0:
                for k in range(len(P)):
                    # if em[i][j] == (P[k]-255):
                    #     img_hide3[i][j] +=secret[i][j]
                    #     ec3-=1
                    if P[k] > Z[k] and ec3 > 0:
                        # if em[index] == (P[k] - 255):
                        if em[i][j] == (P[k] - 255) and ec3 > 0:
                            img_hide3[i][j] -= secret[i][j]
                            ec3 -= 1
                        if em[i][j] > (Z[k] - 255) and em[i][j] < (P[k] - 255):
                        # if em[index]  > (Z[k] - 255) and em[index]  < (P[k] - 255):
                            img_hide3[i][j] -= 1
                            # index += 1
                    if P[k] < Z[k] and ec3 > 0:
                        if em[i][j] == (P[k] - 255):
                        # if em[index]  == (P[k] - 255):
                            img_hide3[i][j] += secret[i][j]
                            ec3 -= 1
                        if em[i][j] < (Z[k] - 255) and em[i][j] > (P[k] - 255):
                        # if em[index] < (Z[k] - 255) and em[index]> (P[k] - 255):
                            img_hide3[i][j] += 1
                index+=1
            if np.mod(i + j, 2) == 1 and flag3 == 1 and ec3 > 0:
                for k in range(len(P)):
                    # if em[i][j] == (P[k] - 255):
                    #     img_hide3[i][j] += secret[i][j]
                    #     ec3 -=1
                    if P[k] > Z[k] and ec3 > 0:
                        # if em[index] == (P[k] - 255):
                        if em[i][j] == (P[k] - 255) and ec3 > 0:
                            img_hide3[i][j] -= secret[i][j]
                            ec3 -= 1
                        if em[i][j] > (Z[k] - 255) and em[i][j] < (P[k] - 255):
                        # if em[index] > (Z[k] - 255) and em[index] < (P[k] - 255):
                            img_hide3[i][j] -= 1
                            # index += 1
                    if P[k] < Z[k] and ec3 > 0:
                        if em[i][j] == (P[k] - 255) and ec3 > 0:
                        # if em[index] == (P[k] - 255):
                            img_hide3[i][j] += secret[i][j]
                            ec3 -= 1
                        if em[i][j] < (Z[k] - 255) and em[i][j] > (P[k] - 255):
                        # if em[index] < (Z[k] - 255) and em[index] > (P[k] - 255):
                            img_hide3[i][j] += 1
                index += 1
    return img_hide3, ec3


def plt_make(y_size,x_size):
    x = x_size
    y = y_size
    plt.plot(x, y, '-r', label='psnr')
    plt.plot(x, y, '.b')
    plt.title('lena psnr')
    plt.xlabel('ec', color='#1C2833')
    plt.ylabel('psnr', color='#1C2833')
    plt.legend(loc='upper left')
    plt.grid()
    plt.show()

def getModifyMap(P,Z):
    P = P+255
    Z = Z+255
    P = np.sort(P)
    Z = np.sort(Z)
    for i in range(len(P)):
        for j in range(len(P)):
            if P[i] < Z[i]:
                if P[i] + 1 <= Z[i] - 1:
                    if P[j]>P[i] and P[j]<Z[i]:
                        P[j]+=1
                        Z[j]+=1

            else:
                if P[i] - 1 >= Z[i] + 1:
                    if P[j] < P[i] and P[j] > Z[i]:
                        P[j] += 1
                        Z[j]-=1
    P =P-255
    Z = Z-255
    return  P,Z


if __name__ == '__main__':

    x_size  =np.array([0.025,0.05,0.1,0.2,0.3,0.4,0.5])
    y_size = []
    for i in range(len(x_size)):
        print('第{}轮插入'.format(i))
        flag = 0
        # ecfinal = 10000
        ecfinal = np.ceil(x_size[i] * 512 * 512)
        k = 16
        Img = Load_img('test', flag)
        img = Img[0]
        imgs = copy.deepcopy(img)
        row, col = imgs.shape
        #单次嵌入
        imgs,ec = danci(imgs,flag,ecfinal)
        img_hide = copy.deepcopy(imgs)
        psnr = cal_psnr(img, img_hide, ec)
        print('psnr',psnr)
        print('ec',ec)
        if ec == 0:
            print('---------------第二次插入--------')
            flag = 1
            new_img, ec = danci(img_hide, flag, ecfinal)
            psnr = cal_psnr(img, img_hide, ec)
            y_size.append(psnr)
            print('psnr', psnr)
            print('ec', ec)
    print(y_size,x_size)
    pd.DataFrame(y_size).to_csv('./psnr.csv')
    plt_make(y_size,x_size)
    print(y_size,x_size)
    # flag = 0
    # # ecfinal = 10000
    # ecfinal = np.ceil(0.1 * 512 * 512)
    # k = 16
    # Img = Load_img('test', flag)
    # img = Img[0]
    # imgs = copy.deepcopy(img)
    # row, col = imgs.shape
    # # 单次嵌入
    # print('ecfinal',ecfinal)
    # imgs, ec = danci(imgs, flag, ecfinal)
    # img_hide = copy.deepcopy(imgs)
    # psnr = cal_psnr(img, img_hide, ec)
    # print('psnr', psnr)
    # print('ec', ec)
    # if ec == 0:
    #     print('---------------第二次插入--------')
    #     flag = 1
    #     new_img, ec = danci(img_hide, flag, ecfinal)
    #     psnr = cal_psnr(img, img_hide, ec)
    #     psnr = cal_psnr(img, new_img, ec)
    #     # y_size.append(psnr)
    #     print('psnr', psnr)
    #     print('ec', ec)